from neo4j import GraphDatabase, exceptions
from config import NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD

class Neo4jUtils:
    def __init__(self):
        self.driver = GraphDatabase.driver(
            NEO4J_URI,
            auth=(NEO4J_USER, NEO4J_PASSWORD)
        )
    
    def close(self):
        """关闭数据库连接"""
        if self.driver:
            self.driver.close()
    
    def search_nodes(self, query_string):
        """
        根据查询字符串搜索节点及其上下游关系
        :param query_string: 搜索字符串
        :return: 节点及其关系信息
        """
        try:
            with self.driver.session() as session:
                # 搜索包含查询字符串的节点
                result = session.run("""
                    MATCH (n)
                    WHERE toLower(n.name) CONTAINS toLower($query) OR toLower(n.content) CONTAINS toLower($query)
                    OPTIONAL MATCH (n)-[r]->(m)  // 下游关系
                    OPTIONAL MATCH (p)-[s]->(n)  // 上游关系
                    RETURN n, collect(DISTINCT {relation: type(r), node: m}) as downstream,
                           collect(DISTINCT {relation: type(s), node: p}) as upstream
                """, query=query_string)
                
                nodes = []
                for record in result:
                    node = record["n"]
                    node_data = {
                        "id": node.id,
                        "labels": list(node.labels),
                        "properties": dict(node)
                    }
                    
                    downstream = []
                    for d in record["downstream"]:
                        if d["node"]:
                            downstream.append({
                                "relation": d["relation"],
                                "node": {
                                    "id": d["node"].id,
                                    "labels": list(d["node"].labels),
                                    "properties": dict(d["node"])
                                }
                            })
                    
                    upstream = []
                    for u in record["upstream"]:
                        if u["node"]:
                            upstream.append({
                                "relation": u["relation"],
                                "node": {
                                    "id": u["node"].id,
                                    "labels": list(u["node"].labels),
                                    "properties": dict(u["node"])
                                }
                            })
                    
                    nodes.append({
                        "node": node_data,
                        "downstream": downstream,
                        "upstream": upstream
                    })
                
                return nodes
        except exceptions.Neo4jError as e:
            print(f"Neo4j error: {e}")
            return []
    
    def update_node(self, node_id, properties):
        """
        更新节点属性
        :param node_id: 节点ID
        :param properties: 要更新的属性字典
        :return: 更新结果
        """
        try:
            with self.driver.session() as session:
                # 构建更新语句
                set_clause = ", ".join([f"n.{key} = ${key}" for key in properties.keys()])
                query = f"MATCH (n) WHERE id(n) = $node_id SET {set_clause} RETURN n"
                
                result = session.run(query, node_id=node_id, **properties)
                updated_node = result.single()
                
                if updated_node:
                    node = updated_node["n"]
                    return {
                        "success": True,
                        "node": {
                            "id": node.id,
                            "labels": list(node.labels),
                            "properties": dict(node)
                        }
                    }
                else:
                    return {"success": False, "message": "Node not found"}
        except exceptions.Neo4jError as e:
            print(f"Neo4j error: {e}")
            return {"success": False, "message": str(e)}
